from flask import Flask, request, jsonify
# from flask_cors import CORS
import torch
from transformers import BertTokenizer, BertForSequenceClassification

app = Flask(__name__)
# CORS(app)  # 允许跨域请求

# 设备设置
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 加载模型和tokenizer
MODEL_NAME = 'bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2).to(device)
model.load_state_dict(torch.load('../model/best_model.bin'))
model.eval()
print("模型加载完毕！")

def predict(text, max_len=128):
    encoding = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=max_len,
        return_token_type_ids=False,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt',
    )
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    
    logits = outputs.logits
    _, prediction = torch.max(logits, dim=1)
    return '攻击性言论' if prediction.item() == 1 else '正常言论'

@app.route('/predict', methods=['POST'])
def api_predict():
    try:
        data = request.values.get('text')
        if not data:
            return jsonify({"error": "Missing 'text' in request body"}), 400
            
        text = data
        result = predict(text)
        if text == "攻击性言论":
            res = 1
        else:
            res = 0
        return jsonify({"message": result,"result":res})
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)